package com.czk.ssaa2aagentcall.controller;

import com.alibaba.cloud.ai.graph.agent.BaseAgent;
import com.alibaba.cloud.ai.graph.agent.a2a.AgentCardProvider;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.util.List;
import java.util.Map;

/**
 * @author chenzhangkun
 * @ClassName A2ARemoteCallController
 * @description: TODO
 * @date 2025年11月19日
 */
@RestController
@Slf4j
public class A2ARemoteCallController {
    @Resource
    private BaseAgent rootAgent;

    @GetMapping(value = "/a2a/call")
    public Flux<String> stream(@RequestParam("question") String question) throws GraphStateException, GraphRunnerException {
        return rootAgent.stream(Map.of("messages", List.of(new UserMessage(question)))).mapNotNull(output -> {
            log.debug("stream agent invoke : `{}`", output.toString());
            if (output.isSTART() || output.isEND()) {
                return null;
            }
            if (output instanceof StreamingOutput) {
                return ((StreamingOutput) output).chunk();
            }
            return null;
        }).publishOn(Schedulers.parallel());
    }

    @GetMapping("sync")
    public Object sync(@RequestParam("question") String question) throws GraphStateException, GraphRunnerException {
        System.out.println(question);
        return rootAgent.invoke(Map.of("messages", List.of(new UserMessage(question)))).orElseThrow().value("messages")
                .orElseThrow();
    }


}
